# -*- coding: utf-8 -*-
# a test
from scipy import *
from generator import *
from pysw import SWNPs


def simulate_zfc_fc(NP, field, temp_list):
    '''The function to perform a ZFC-FC simulation. The NP should first be initialized.
    The function performs the operation on the NP object.
    The result of the simulation will be generated.
    '''
    temp_list.sort()
    NP.updateField(field)
    ZFC_moment = []
    for temp in temp_list:
        NP.updateTemperature(temp)
        ZFC_moment.append(NP.getMoment())
    ZFC_moment = array(ZFC_moment)
    FC_moment = []
    for temp in temp_list:
        NP.updateTemperature(temp)
        FC_moment.append(NP.getMoment())
    FC_moment = array(FC_moment)
    return ZFC_moment, FC_moment

def fit_zfc_fc_ka(ka, templist, ZFCdata, FCdata, N, theta, phi, volume, magnetisation, field):
    # The ZFC and FC data has to be normalized. Else it will NEVER work!!!!
    # The data is normalized so that the max of all data is 1, and min of all is 0
    FitNP = SWNPs(N, theta, phi, volume, ka, 0, magnetisation, 0.00001)
    templist.sort()
    ZFCfit, FCfit = simulate_zfc_fc(FitNP, field, templist)
    norm_up = max(max(FCfit), max(ZFCfit))
    norm_down = min(min(FCfit), min(ZFCfit))
    #ZFCfit = (ZFCfit-norm_down)/(norm_up-norm_down)
    #FCfit = (FCfit-norm_down)/(norm_up-norm_down)
    ZFCfit = ZFCfit/norm_up
    FCfit = FCfit/norm_up
    err = concatenate((ZFCdata-ZFCfit, FCdata-FCfit))
    print 'estimate error is: %f' %(err**2).sum()
    return err

def fit_zfc_fc_ka_mag(x0 , templist, ZFCdata, FCdata, N, theta, phi, volume, field):
    # The ZFC and FC data has to be normalized. Else it will NEVER work!!!!
    # The data is normalized so that the max of all data is 1, and min of all is 0
    ka, magnetisation = x0
    FitNP = SWNPs(N, theta, phi, volume, ka, 0, magnetisation, 0.00001)
    templist.sort()
    ZFCfit, FCfit = simulate_zfc_fc(FitNP, field, templist)
    norm_up = max(max(FCfit), max(ZFCfit))
    norm_down = min(min(FCfit), min(ZFCfit))
    ZFCfit = (ZFCfit-norm_down)/(norm_up-norm_down)
    FCfit = (FCfit-norm_down)/(norm_up-norm_down)
    #ZFCfit = ZFCfit/norm_up
    #FCfit = FCfit/norm_up
    err = concatenate((ZFCdata-ZFCfit, FCdata-FCfit))
    print 'use ka=%f ms=%f estimate error is: %f' %(ka,magnetisation,(err**2).sum())
    return err

def fit_zfc_fc_ka_size(x0 , templist, ZFCdata, FCdata, N, theta, phi, magnetisation, field):
    # fit the ZFCFC with Ka and size distribution
    ka, mu, sigma = x0
    def generate_lognorm(n=None):
        '''generate the volume distribution for a spherical particle with log-normal dist.
        Not the average and variance of the log-normal distribution is NOT the mu and sigma factor here.
        One should calculate the value again.
        
        Use the same hack as the above function, in order to reduce numeric error
        '''
        #mu = 2.58041
        #sigma = 0.121255
        #global mu
        #global sigma
        m = n/2
        dia = random.lognormal(mu,sigma,m)*1e-9
        volume = pi*dia**3/6
        volume = concatenate((volume, volume))
        return volume
    
    FitNP = SWNPs(N, theta, phi, generate_lognorm, ka, 0, magnetisation, 0.00001)
    templist.sort()
    ZFCfit, FCfit = simulate_zfc_fc(FitNP, field, templist)
    norm_up = max(max(FCfit), max(ZFCfit))
    norm_down = min(min(FCfit), min(ZFCfit))
    ZFCfit = (ZFCfit-norm_down)/(norm_up-norm_down)
    FCfit = (FCfit-norm_down)/(norm_up-norm_down)
    #ZFCfit = ZFCfit/norm_up
    #FCfit = FCfit/norm_up
    err = concatenate((ZFCdata-ZFCfit, FCdata-FCfit))
    print 'use ka=%f mu=%f estimate error is: %f' %(ka,mu,(err**2).sum())
    return err


def fit_zfc_fc_ka_size_para(x0 , templist, ZFCdata, FCdata, N, theta, phi, magnetisation, field):
    # fit the ZFCFC with Ka and size distribution
    ka, mu, sigma, para = x0
    if sigma < 0: sigma = 1e-5
    if mu < 0: mu = 1e-5
    def generate_lognorm(n=None):
        '''generate the volume distribution for a spherical particle with log-normal dist.
        Not the average and variance of the log-normal distribution is NOT the mu and sigma factor here.
        One should calculate the value again.
        
        Use the same hack as the above function, in order to reduce numeric error
        '''
        #mu = 2.58041
        #sigma = 0.121255
        #global mu
        #global sigma
        m = n/2
        dia = random.lognormal(mu,sigma,m)*1e-9
        volume = pi*dia**3/6
        volume = concatenate((volume, volume))
        return volume
    
    FitNP = SWNPs(N, theta, phi, generate_lognorm, ka, 0, magnetisation, 0.00001)
    templist.sort()
    ZFCfit, FCfit = simulate_zfc_fc(FitNP, field, templist)
    ZFCfit = ZFCfit + para/templist
    FCfit = FCfit + para/templist
    norm_up = max(max(FCfit), max(ZFCfit))
    norm_down = min(min(FCfit), min(ZFCfit))
    ZFCfit = (ZFCfit-norm_down)/(norm_up-norm_down)
    FCfit = (FCfit-norm_down)/(norm_up-norm_down)
    #ZFCfit = ZFCfit/norm_up
    #FCfit = FCfit/norm_up
    err = concatenate((ZFCdata-ZFCfit, FCdata-FCfit))
    print 'use ka=%f mu=%f simga=%f para=%f estimate error is: %f' %(ka,mu,sigma,para,(err**2).sum())
    return err



if __name__ == '__main__':
    # For test only
    #TestNP = SWNPs(5000, generate_theta, generate_phi, generate_volume, 1e4, 0, 1700e3, 0.00001)
    #temp_list = arange(5,300,5)
    #zfc, fc = simulate_zfc_fc(TestNP, 1e-3, temp_list)
    #out = array([temp_list, zfc, fc]).transpose()
    #savetxt('testzfcfc', out)
    
    zfcfile = 'ZFC-set1.dat'
    fcfile = 'FC-set1.dat'
    #Ka_start = 3e4
    #Ka_start = 5.4e4
    Ka_start = 40000
    #mu_start = 2.58041
    #sigma_start = 0.121255
    N = 10000
    #magnetisation = 1700e3
    #magnetisation = 169815
    magnetisation = 170000
    #para_start = 5e-16
    field = 0.01
    ZFCdatas = loadtxt(zfcfile)
    FCdatas = loadtxt(fcfile)
    templist = ZFCdatas[:,0]
    ZFCdata = ZFCdatas[:,1]
    FCdata = FCdatas[:,1]
    # normalize the data
    norm_up = max(max(FCdata), max(ZFCdata))
    norm_down = min(min(FCdata), min(ZFCdata))
    ZFCdata = (ZFCdata-norm_down)/(norm_up-norm_down)
    FCdata = (FCdata-norm_down)/(norm_up-norm_down)
    #ZFCdata = ZFCdata / norm_up
    #FCdata = FCdata / norm_up
    #start the simualtion
    #res = optimize.leastsq(fit_zfc_fc_ka, Ka_start,args=(templist, ZFCdata, FCdata, N, generate_theta, generate_phi, generate_volume, magnetisation, field))
    res = optimize.leastsq(fit_zfc_fc_ka_mag,array([Ka_start, magnetisation]),args=(templist, ZFCdata, FCdata, N, generate_theta, generate_phi, generate_volume, field))
    #res = optimize.leastsq(fit_zfc_fc_ka_size,array([Ka_start, mu_start, sigma_start]),args=(templist, ZFCdata, FCdata, N, generate_theta, generate_phi, magnetisation, field))
    #res = optimize.leastsq(fit_zfc_fc_ka_size_para,array([Ka_start, mu_start, sigma_start, para_start]),args=(templist, ZFCdata, FCdata, N, generate_theta, generate_phi, magnetisation, field))
    
    print res
    Ka = res[0][0]
    magnetisation = res[0][1]
    #mu = res[0][1]
    #sigma = res[0][2]
    #para = res[0][3]
    
    def generate_lognorm(n=None):
        '''generate the volume distribution for a spherical particle with log-normal dist.
        Not the average and variance of the log-normal distribution is NOT the mu and sigma factor here.
        One should calculate the value again.
        
        Use the same hack as the above function, in order to reduce numeric error
        '''
        #mu = 2.58041
        #sigma = 0.121255
        #global mu
        #global sigma
        m = n/2
        #if sigma < 0: sigma = 1e-5
        #if mu < 0: mu = 1e-5
        dia = random.lognormal(mu,sigma,m)*1e-9
        volume = pi*dia**3/6
        volume = concatenate((volume, volume))
        return volume
    
    FitNP = SWNPs(N, generate_theta, generate_phi, generate_volume, Ka, 0, magnetisation, 0.00001)
    ZFCfit, FCfit = simulate_zfc_fc(FitNP, field, templist)
    #ZFCfit = ZFCfit + para/templist
    #FCfit = FCfit + para/templist
    ZFCfit = array([templist,ZFCfit]).transpose()
    FCfit = array([templist,FCfit]).transpose()
    savetxt('ZFCfit',ZFCfit)
    savetxt('FCfit',FCfit)

    
